#include "linux/string.h"
#include "linux/printk.h"
#include "linux/common.h"
#include "linux/types.h"
#include "linux/kernel.h"

#define MAX_REGIONS 128

// 内存区域描述符
struct memblock_region
{
    unsigned long base;
    size_t size;
};

// memblock 分配器结构
struct memblock
{
    struct memblock_region *regions;
    unsigned long count;
    unsigned long max;
};

// 全局 memblock
static struct memblock memory = {0};
static struct memblock reserved = {0};

// 内存区域数组（静态分配）
static struct memblock_region memory_regions[MAX_REGIONS];
static struct memblock_region reserved_regions[MAX_REGIONS];


// 初始化 memblock 系统
void memblock_init(void)
{
    // 初始化 memory memblock
    memory.regions = memory_regions;
    memory.max = MAX_REGIONS;
    memory.count = 0;

    // 初始化 reserved memblock
    reserved.regions = reserved_regions;
    reserved.max = MAX_REGIONS;
    reserved.count = 0;

}

// 内部函数：在 memblock 中添加区域
static int memblock_add_region(struct memblock *mb, unsigned long base, size_t size)
{
    if (mb->count >= mb->max)
        return -1; // 区域已满

    // 添加到末尾
    mb->regions[mb->count].base = base;
    mb->regions[mb->count].size = size;
    mb->count++;

    return 0;
}

// 内部函数：从 memblock 中移除区域
static int memblock_remove_region(struct memblock *mb, unsigned long base, size_t size)
{
    unsigned int i;
    
    for (i = 0; i < mb->count; i++)
    {
        if (mb->regions[i].base == base && mb->regions[i].size == size)
        {
            // 移除区域：将最后一个区域移动到当前位置
            if (i < mb->count - 1)
            {
                mb->regions[i] = mb->regions[mb->count - 1];
            }
            mb->count--;
            return 0;
        }
    }

    return -1; // 未找到匹配区域
}

// 添加内存区域
int memblock_add(unsigned long base, size_t size)
{
    return memblock_add_region(&memory, base, size);
}

// 移除内存区域
int memblock_remove(unsigned long base, size_t size)
{
    return memblock_remove_region(&memory, base, size);
}

// 保留内存区域
int memblock_reserve(unsigned long base, size_t size)
{
    return memblock_add_region(&reserved, base, size);
}

// 释放内存区域
int memblock_free(unsigned long base, size_t size)
{
    return memblock_remove_region(&reserved, base, size);
}

// 检查地址是否在保留区域内
static int is_reserved(unsigned long addr, size_t size)
{
    unsigned long i;
    size_t r_size;
    unsigned long r_base;
    
    for (i = 0; i < reserved.count; i++)
    {
        r_base = reserved.regions[i].base;
        r_size = reserved.regions[i].size;

        // 检查是否有重叠
        if (addr < r_base + r_size && addr + size > r_base)
        {
            return 1; // 有重叠，已保留
        }
    }

    return 0; // 未保留
}

// 分配内存
void *memblock_alloc(size_t size, size_t align)
{
    // 计算对齐后的地址
    unsigned long aligned_addr;
    unsigned long base;
    size_t total_size;
    
    // 确保有内存区域
    if (memory.count == 0)
        return NULL;

    // 获取内存区域的基址和大小
    base = memory.regions[0].base;
    total_size = memory.regions[0].size;

    // 尝试从内存起始处重新分配
    aligned_addr = (base + align - 1) & ~(align - 1);
    if (aligned_addr + size > base + total_size)
        return NULL; // 内存不足

    // 检查是否与保留区域重叠
    while (is_reserved(aligned_addr, size))
    {
        // 跳过已保留的区域
        aligned_addr += align;

        // 检查是否超出内存范围
        if (aligned_addr + size > base + total_size)
        {
            return NULL; // 内存不足
        }
    }

    // 保留这块内存
    if (memblock_reserve(aligned_addr, size) != 0)
        return NULL; // 保留失败

    return (void *)aligned_addr;
}

int __next_reserved_mem_region(int idx, phys_addr_t *out_start, phys_addr_t *out_size)
{
    struct memblock *type = &reserved;

    if (idx < type->count)
    {
        struct memblock_region *r = &type->regions[idx];
        phys_addr_t base = r->base;
        phys_addr_t size = r->size;

        if (out_start)
            *out_start = base;
        if (out_size)
            *out_size = size;

        return  0;
    }

    return -1;
}

// 获取内存信息
void memblock_get_info(unsigned long *total_memory, unsigned long *used_memory)
{
    unsigned long i;

    *total_memory = 0;
    *used_memory = 0;

    // 计算总内存
    for (i = 0; i < memory.count; i++)
        *total_memory += memory.regions[i].size;

    // 计算已使用内存
    for (i = 0; i < reserved.count; i++)
        *used_memory += reserved.regions[i].size;
}

// 调试函数：打印内存状态
void memblock_dump_all(void)
{
    unsigned int i;
    
    printk("Memory regions:\n");
    for (i = 0; i < memory.count; i++)
    {
        printk("  [%d] 0x%lx - 0x%lx (size: 0x%lx)\n",
               i, memory.regions[i].base,
               memory.regions[i].base + memory.regions[i].size - 1,
               memory.regions[i].size);
    }

    printk("Reserved regions:\n");
    for (i = 0; i < reserved.count; i++)
    {
        printk("  [%d] 0x%lx - 0x%lx (size: 0x%lx)\n",
               i, 
               reserved.regions[i].base,
               reserved.regions[i].base + reserved.regions[i].size,
               reserved.regions[i].size);
    }

}